package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"syscall"
	"time"

	"github.com/google/uuid"
)

type TaskRequest struct {
	TaskType     string            `json:"task_type"`              // "command" | "script"
	Command      string            `json:"command,omitempty"`
	ScriptURL    string            `json:"script_url,omitempty"`
	Interpreter  string            `json:"interpreter"`            // e.g. "/bin/bash"
	Env          map[string]string `json:"env,omitempty"`
	Mode         string            `json:"mode"`                   // "sync","async","auto"
	TimeoutSec   int               `json:"timeout"`                // auto同步窗口 秒
	SyncTimeout  int               `json:"sync_timeout,omitempty"` // sync最长超时 秒，默认300
	AsyncTimeout int               `json:"async_timeout,omitempty"`// async最长 秒，默认24*3600
	CallbackURL  string            `json:"callback"`               // 回调url
	TaskID       string            `json:"task_id"`                // 唯一ID
}

type TaskResult struct {
	TaskID   string `json:"task_id"`
	Stdout   string `json:"stdout"`
	Stderr   string `json:"stderr"`
	ExitCode int    `json:"exit_code"`
	Timeout  bool   `json:"timeout,omitempty"`
}

// 用于异步管理正在执行的进程
type runningTaskInfo struct {
	Cmd    *exec.Cmd
	Cancel context.CancelFunc
	Done   chan struct{}
	Once   sync.Once
}

// 全局存储(可选，用于并发管理/手动cancel简单设计即可)
var (
	tasks      = make(map[string]*TaskResult)
	runningMap = make(map[string]*runningTaskInfo)
	tasksMu    sync.Mutex
	defaultSyncTimeout  = 300
	defaultAsyncTimeout = 86400
)

func main() {
	sock := "/var/run/supervisor.sock"
	os.RemoveAll(sock)
	l, err := net.Listen("unix", sock)
	if err != nil {
		log.Fatal("listen unix socket:", err)
	}
	http.HandleFunc("/run", handleRun)
	log.Println("supervisor-server listening on", sock)
	log.Fatal(http.Serve(l, nil))
}

func handleRun(w http.ResponseWriter, r *http.Request) {
	var req TaskRequest
	b, _ := ioutil.ReadAll(r.Body)
	if err := json.Unmarshal(b, &req); err != nil {
		http.Error(w, err.Error(), 400)
		return
	}

	if req.SyncTimeout < 1 {
		req.SyncTimeout = defaultSyncTimeout
	}
	if req.AsyncTimeout < 1 {
		req.AsyncTimeout = defaultAsyncTimeout
	}
	if req.TimeoutSec < 1 {
		req.TimeoutSec = 1
	}
	if req.Interpreter == "" {
		req.Interpreter = "/bin/bash"
	}
	if req.TaskID == "" {
		req.TaskID = uuid.New().String()
	}

	tr := &TaskResult{TaskID: req.TaskID}
	tasksMu.Lock()
	tasks[req.TaskID] = tr
	tasksMu.Unlock()

	switch req.Mode {
	case "sync":
		doSync(w, req, tr)
	case "async":
		w.Write([]byte(fmt.Sprintf(`{"task_id":"%s"}`, req.TaskID)))
		go doAsync(req, tr)
	case "auto":
		doAuto(w, req, tr)
	default:
		http.Error(w, "mode must be sync|async|auto", 400)
	}
}

// ---同步模式---
func doSync(w http.ResponseWriter, req TaskRequest, tr *TaskResult) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(req.SyncTimeout)*time.Second)
	defer cancel()
	exit, timeout := runTaskWithCtx(ctx, req, tr)
	tr.Timeout = timeout
	tr.ExitCode = exit
	payload, _ := json.Marshal(tr)
	w.Write(payload)
}

// ---异步模式---
func doAsync(req TaskRequest, tr *TaskResult) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(req.AsyncTimeout)*time.Second)
	defer cancel()
	exit, timeout := runTaskWithCtx(ctx, req, tr)
	tr.Timeout = timeout
	tr.ExitCode = exit
	if req.CallbackURL != "" {
		payload, _ := json.Marshal(tr)
		http.Post(req.CallbackURL, "application/json", bytes.NewReader(payload))
	}
}

// ---修正后自动分流模式---
func doAuto(w http.ResponseWriter, req TaskRequest, tr *TaskResult) {
	autoTimeout := time.Duration(req.TimeoutSec) * time.Second
	asyncTimeout := time.Duration(req.AsyncTimeout) * time.Second

	done := make(chan struct{})
	var exitcode int
	var timeoutFlag bool

	// 用一个完全独立的ctx, asyncTimeout控制，确保不会受到主流程影响
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), asyncTimeout)
		defer cancel()
		exitcode, timeoutFlag = runTaskWithCtx(ctx, req, tr)
		close(done)
	}()

	select {
	case <-done:
		tr.Timeout = timeoutFlag
		tr.ExitCode = exitcode
		payload, _ := json.Marshal(tr)
		w.Write(payload)
	case <-time.After(autoTimeout):
		// 进入异步态，前端返回，仅等后台任务自己归档
		w.Write([]byte(fmt.Sprintf(`{"task_id":"%s"}`, req.TaskID)))
		go func() {
			<-done
			// 或可加'二次保护'超时select, 理论用不到
			tr.Timeout = timeoutFlag
			tr.ExitCode = exitcode
			if req.CallbackURL != "" {
				payload, _ := json.Marshal(tr)
				http.Post(req.CallbackURL, "application/json", bytes.NewReader(payload))
			}
		}()
	}
}

func runTaskWithCtx(ctx context.Context, req TaskRequest, tr *TaskResult) (exitCode int, timeout bool) {
	var cmd *exec.Cmd

	if req.TaskType == "script" {
		tmpDir := "/tmp/supervisor-scripts"
		os.MkdirAll(tmpDir, 0755)
		ext := filepath.Ext(req.ScriptURL)
		if ext == "" {
			ext = ".sh"
		}
		tmpFile := filepath.Join(tmpDir, req.TaskID+ext)
		resp, err := http.Get(req.ScriptURL)
		if err != nil {
			tr.Stderr = err.Error()
			return -1, false
		}
		defer resp.Body.Close()
		data, _ := io.ReadAll(resp.Body)
		if err := ioutil.WriteFile(tmpFile, data, 0700); err != nil {
			tr.Stderr = err.Error()
			return -1, false
		}
		cmd = exec.CommandContext(ctx, req.Interpreter, tmpFile)
		defer os.Remove(tmpFile)
	} else {
		cmd = exec.CommandContext(ctx, req.Interpreter, "-c", req.Command)
	}

	env := os.Environ()
	for k, v := range req.Env {
		env = append(env, k+"="+v)
	}
	cmd.Env = env

	var outBuf, errBuf bytes.Buffer
	cmd.Stdout = &outBuf
	cmd.Stderr = &errBuf

	rt := &runningTaskInfo{
		Cmd:    cmd,
		Done:   make(chan struct{}),
	}
	ctx2, cancel := context.WithCancel(ctx)
	rt.Cancel = cancel

	tasksMu.Lock()
	runningMap[req.TaskID] = rt
	tasksMu.Unlock()

	defer func() {
		rt.Once.Do(func() { close(rt.Done) })
		tasksMu.Lock()
		delete(runningMap, req.TaskID)
		tasksMu.Unlock()
	}()

	err := cmd.Start()
	if err != nil {
		tr.Stderr = err.Error()
		return -1, false
	}
	errch := make(chan error, 1)
	go func() {
		errch <- cmd.Wait()
	}()

	select {
	case <-ctx2.Done():
		killCmd(cmd)
		// （注意，只覆盖一次，不叠加）
		if tr.Stderr == "" {
			tr.Stderr = "[killed by agent: timeout] " + ctx2.Err().Error()
		}
		return -1, true
	case err := <-errch:
		tr.Stdout = outBuf.String()
		tr.Stderr = errBuf.String()
		if err != nil {
			if exiterr, ok := err.(*exec.ExitError); ok {
				if ws, ok2 := exiterr.Sys().(syscall.WaitStatus); ok2 {
					return ws.ExitStatus(), false
				}
			}
			return -1, false
		}
		if cmd.ProcessState != nil {
			return cmd.ProcessState.ExitCode(), false
		}
		return 0, false
	}
}

func killTask(taskID string) {
	tasksMu.Lock()
	rt, ok := runningMap[taskID]
	tasksMu.Unlock()
	if ok && rt != nil && rt.Cmd != nil && rt.Cmd.Process != nil {
		killCmd(rt.Cmd)
		rt.Once.Do(func() { close(rt.Done) })
	}
}

func killCmd(cmd *exec.Cmd) {
	if cmd.Process == nil {
		return
	}
	_ = cmd.Process.Kill()
}



